package com.sub.spark.sql.udf;

import org.apache.spark.sql.Encoder;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.expressions.Aggregator;

import java.io.Serializable;

import static org.apache.spark.sql.functions.udaf;

/**
 * @ClassName SparkSqlUDAF
 * @Description: UDAF：输入多行，返回一行（聚合函数）。通常和groupBy一起使用，如果直接使用UDAF函数，默认将所有的数据合并在一起。
 * @Author Submerge.
 * @Since 2025/5/24 13:48
 * @Version 1.0
 */
public class SparkSqlUDAF {

    public static void main(String[] args) {

        SparkSession spark = SparkSession.builder().appName("sub-spark-sql-udaf").master("local[2]").getOrCreate();

        spark.read().json("data/demo/spark/user.json").createOrReplaceTempView("user");

        // 注册需要导入依赖 import static org.apache.spark.sql.functions.udaf;
        spark.udf().register("subAvg",udaf(new MyAvg(), Encoders.LONG()));

        spark.sql("select subAvg(age) newAge from user").show();

        //4. 关闭sparkSession
        spark.close();
    }

    /**
     * 定义一个缓冲区，用于保存中间结果
     */
    public static class Buffer implements Serializable {

        private Long sum;
        private Long count;
        public Buffer() {
        }
        public Buffer(Long sum, Long count) {
            this.sum = sum;
            this.count = count;
        }
        public Long getSum() {
            return sum;
        }
        public void setSum(Long sum) {
            this.sum = sum;
        }
        public Long getCount() {
            return count;
        }
        public void setCount(Long count) {
            this.count = count;
        }
    }


    /**
     * 自定义聚合函数类 继承Aggregator，泛型为<输入类型，中间缓冲区类型，输出类型>，实现方法
     */
    public static class MyAvg extends Aggregator<Long,Buffer,Double> {

        /**
         * 初始化缓冲区
         * @return
         */
        @Override
        public Buffer zero() {
            return new Buffer(0L,0L);
        }

        /**
         * 将输入和缓冲区聚合
         * @param b
         * @param a
         * @return
         */
        @Override
        public Buffer reduce(Buffer b, Long a) {
            b.setSum(b.getSum() + a);
            b.setCount(b.getCount() + 1);
            return b;
        }

        /**
         * 将多个缓冲区聚合
         * @param b1
         * @param b2
         * @return
         */
        @Override
        public Buffer merge(Buffer b1, Buffer b2) {
            b1.setSum(b1.getSum() + b2.getSum());
            b1.setCount(b1.getCount() + b2.getCount());
            return b1;
        }

        /**
         * 计算最终结果
         * @param reduction
         * @return
         */
        @Override
        public Double finish(Buffer reduction) {
            return reduction.getSum().doubleValue() / reduction.getCount();
        }

        /**
         * 指定中间缓冲区的编码器
         * @return
         */
        @Override
        public Encoder<Buffer> bufferEncoder() {
            // 可以用kryo进行优化
            return Encoders.kryo(Buffer.class);
        }

        /**
         * 指定输出类型的编码器
         * @return
         */
        @Override
        public Encoder<Double> outputEncoder() {
            return Encoders.DOUBLE();
        }
    }
}